import os
try:
    from .cvio import cvio
except:
    from cvio import cvio
import random
import shutil

def sample_dataset(cutsrc, dst, maxNumPerCls=10000):
    imglist = cvio.load_image_list(cutsrc, recursive=True, silent=False)
    n = len(imglist)
    print('Loaded %d images.' % n)
    if n <= 0: return
    if not os.path.exists(dst):
        os.makedirs(dst)
    cls2imgs = {}
    for i, img in enumerate(imglist, 1):
        label = img.split(os.sep)[-2]
        if not label in cls2imgs:
            cls2imgs[label] = [img]
        else:
            cls2imgs[label] += [img]
    
    m = max(1, len(cls2imgs))
    for i, (label, imgs) in enumerate(cls2imgs.items(), 1):
        c = len(imgs)
        print('[%d/%d] %.2f%% %s' % (i, m, i / m * 100, label))
        x = int(maxNumPerCls * 1.5)
        # if c < x and x < 100:
        #     print('%s框数（%d）过少，不予采样.' % (label, c))
        #     continue
        imgs = random.sample(imgs, min(maxNumPerCls, c))
        saveroot = os.path.join(dst, label)
        if not os.path.exists(saveroot):
            os.makedirs(saveroot)
        for img in imgs:
            shutil.move(img, os.path.join(saveroot, os.path.basename(img)))


def sample_det_dataset(src, dst, ann_type='json', ntest=300):
    img_ann_list = cvio.load_img_ann_list(src, ann_type=ann_type, silent=False)
    n = len(img_ann_list)
    print("Found %d image&annotation pairs from %s." % (n, src))
    if n == 0:
        return
    if not os.path.exists(dst):
        os.makedirs(dst)
    ntest = min(ntest, n)
    test = random.sample(img_ann_list, ntest)
    for i, (img, ann) in enumerate(test, 1):
        bimg = os.path.basename(img)
        print('[%d/%d] %s' % (i, ntest, bimg))
        shutil.move(img, os.path.join(dst, bimg))
        shutil.move(ann, os.path.join(dst, os.path.basename(ann)))
